import cv2
import time
import numpy as np
from pre import preProcess


def region_limit(img):
    rows, cols = img.shape[0], img.shape[1]
    mask = np.zeros((rows, cols), np.uint8)
    mask_height, mask_width = mask.shape[:2]
    # 限制在中心区域
    center_x, center_y = mask_width // 2, mask_height // 2
    # 区域大小为100*100
    region_size = 150
    cv2.rectangle(mask, (center_x - region_size // 2, center_y - region_size // 2),
                  (center_x + region_size // 2, center_y + region_size // 2), 255, -1)
    return mask

def get_features(img):
    start = time.time()
    # 创建SIFT提取器
    sift = cv2.SIFT_create()
    mask = region_limit(img)
    # 检测关键点和计算描述符
    keypoints, descriptors = sift.detectAndCompute(img, mask)
    end = time.time()
    time_scope = end - start
    keypoints_img = cv2.drawKeypoints(img, keypoints, mask)
    return keypoints, descriptors, time_scope, keypoints_img

def get_feature_orb(img):
    start = time.time()
    # 创建ORB提取器
    orb = cv2.ORB_create()
    mask = region_limit(img)
    # 检测关键点和计算描述符
    keypoints, descriptors = orb.detectAndCompute(img, mask)
    end = time.time()
    time_scope = end - start
    keypoints_img = cv2.drawKeypoints(img, keypoints, mask, flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)
    return keypoints, descriptors, time_scope, keypoints_img

def get_feature_akaze(img):
    start = time.time()
    # 创建AKAZE提取器
    akaze = cv2.AKAZE_create()
    mask = region_limit(img)
    # 检测关键点和计算描述符
    keypoints, descriptors = akaze.detectAndCompute(img, mask)
    end = time.time()
    time_scope = end - start
    keypoints_img = cv2.drawKeypoints(img, keypoints, mask, flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)
    return keypoints, descriptors, time_scope, keypoints_img

def get_feature_brisk(img):
    start = time.time()
    # 创建BRISK提取器
    brisk = cv2.BRISK_create()
    mask = region_limit(img)
    # 检测关键点和计算描述符
    keypoints, descriptors = brisk.detectAndCompute(img, mask)
    end = time.time()
    time_scope = end - start
    keypoints_img = cv2.drawKeypoints(img, keypoints, mask, flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)
    return keypoints, descriptors, time_scope, keypoints_img

def show_result(fingerprint_test):
    fingerprint_test = preProcess(fingerprint_test)
    cv2.imshow("Original", cv2.resize(fingerprint_test, None, fx=1, fy=1))
    kp, dp, t, kpi = get_features(fingerprint_test)
    kp_orb, dp_orb, t_orb, kpi_orb = get_feature_orb(fingerprint_test)
    kp_akaze, dp_akaze, t_akaze, kpi_akaze = get_feature_akaze(fingerprint_test)
    kp_brisk, dp_brisk, t_brisk, kpi_brisk = get_feature_brisk(fingerprint_test)
    print(f'SIFT run time:{t}')
    print(f'SIFT feature points num:{len(kp)}')
    print(f'ORB run time:{t_orb}')
    print(f'ORB feature points num:{len(kp_orb)}')
    print(f'AKAZE run time:{t_akaze}')
    print(f'AKAZE feature points num:{len(kp_akaze)}')
    print(f'BRISK run time:{t_brisk}')
    print(f'BRISK feature points num:{len(kp_brisk)}')
    cv2.imshow("SIFT_keypoints", cv2.resize(kpi, None, fx=1, fy=1))
    cv2.imshow("ORB_keypoints", cv2.resize(kpi_orb, None, fx=1, fy=1))
    cv2.imshow("AKAZE_keypoints", cv2.resize(kpi_akaze, None, fx=1, fy=1))
    cv2.imshow("BRISK_keypoints", cv2.resize(kpi_brisk, None, fx=1, fy=1))
    cv2.waitKey(0)

if __name__ == "__main__":
    fingerprint_test = cv2.imread("./data/1_2.tif", cv2.IMREAD_GRAYSCALE)
    show_result(fingerprint_test)
    cv2.destroyAllWindows()